import torch
from torch import nn
import numpy as np
from .quanternion import *
import matplotlib.pyplot as plt
from dataset.utils import quat_slerp, pos_lerp


def PLU(x, alpha=0.1, c=1.0):
    relu = nn.ReLU()
    o1 = alpha * (x + c) - c
    o2 = alpha * (x - c) + c
    o3 = x - relu(x - o2)
    o4 = relu(o1 - o3) + o3
    return o4


def cal_ADE_metric(pred, gt):
    """
    Average displacement error
    :param pred: pred global position, TxJx3
    :param gt: gt global position
    :return: ADE
    """
    diff = pred - gt
    ade = np.linalg.norm(diff, ord=2, axis=2).mean()
    return ade


def cal_FDE_metric(pred, gt):
    """
    Final displacement error
    :param pred: pred global position, TxJx3
    :param gt: gt global position
    :return:
    """
    last_diff = pred[-1] - gt[-1]
    fde = np.linalg.norm(last_diff, ord=2, axis=1).mean()
    return fde


def gen_ztta(distance=50, dim=256, basis=10000):
    ztta = np.zeros((1, distance, dim))
    for i in range(distance):
        time_to_arrive = distance - i
        for d in range(dim):
            if dim % 2 == 0:
                ztta[:, i, d] = np.sin(time_to_arrive / np.power(basis, d / dim))
                # ztta[:, i, d] = np.sin(time_to_arrive / np.power(basis, time_to_arrive / dim))  # old model
            else:
                ztta[:, i, d] = np.cos(time_to_arrive / np.power(basis, (d - 1) / dim))
                # ztta[:, i, d] = np.sin(time_to_arrive / np.power(basis, (time_to_arrive - 1) / dim))  # old model

    return torch.from_numpy(ztta.astype(np.float_))  # 1x50x256


def gen_batch_ztta(len_list, distance=50, dim=256, basis=10000):
    # ztta = torch.zeros((len_list.shape[0], distance, dim))
    ztta = np.zeros((len_list.shape[0], distance, dim))
    for i in range(distance):
        time_to_arrive = len_list - i
        for d in range(dim):
            if dim % 2 == 0:
                ztta[:, i, d] = torch.sin(time_to_arrive / torch.pow(basis, time_to_arrive / dim))
                # ztta[:, i, d] = np.sin(time_to_arrive / np.power(basis, time_to_arrive / dim))
            else:
                # ztta[:, i, d] = torch.cos(time_to_arrive / torch.pow(basis, (time_to_arrive - 1) / dim))
                ztta[:, i, d] = np.cos(time_to_arrive / np.power(basis, (time_to_arrive - 1) / dim))

    # return ztta
    return torch.from_numpy(ztta.astype(np.float_))  # Bx50x256


def gen_ztarget(steps, distance=50):
    time_to_arrive = distance - steps - 1
    if time_to_arrive < 5:
        lambda_target = 0
    elif time_to_arrive >= int(distance * 0.6):

        lambda_target = 1
    else:
        lambda_target = (time_to_arrive - 5) / (5 + int(distance * 0.6))

    return lambda_target


def gen_ztarget_by_tta(tta, distance=50):
    time_to_arrive = tta
    if time_to_arrive < 5:
        lambda_target = 0
    elif time_to_arrive >= int(distance * 0.6):
        lambda_target = 1
    else:
        lambda_target = (time_to_arrive - 5) / (5 + int(distance * 0.6))

    return lambda_target


def gen_ratio(steps, distance=50):
    time_to_arrive = distance - steps - 1
    if time_to_arrive < 5:
        lambda_target = 0.7
    elif time_to_arrive >= distance - 5:
        lambda_target = 0
    else:
        lambda_target = (time_to_arrive - 5) / (distance - 10)

    return lambda_target


def gen_target_loss_weight(steps, distance=50):
    time_to_arrive = distance - steps - 1
    if time_to_arrive < 5:
        lambda_target = 1 - time_to_arrive / 5
    else:
        lambda_target = 0

    return lambda_target


def write_to_bvhfile(data, filename):
    fout = open(filename, 'w')
    line_cnt = 0
    for line in open('D:\python\Sketch2Seq\src\example.bvh', 'r'):
        fout.write(line)
        line_cnt += 1
        if line_cnt >= 132:
            break
    fout.write(('Frames: %d\n' % data.shape[0]))
    fout.write('Frame Time: 0.008333\n')
    pose_data = qeuler_np(data[:, 3:].reshape(data.shape[0], -1, 4), order='xyz', use_gpu=False)
    pose_data = pose_data / np.pi * 180.0
    
    # change order
    order = [0, 1, 4, 7, 10, 2, 5, 8, 11, 3, 6, 9, 12, 15, 13, 16, 18, 20, 14, 17, 19, 21]
    pose_data_bvh = np.array([])
    for i in order:
        if i==0:
            pose_data_bvh = pose_data[:,i,:]
        else:
            pose_data_bvh = np.concatenate((pose_data_bvh, pose_data[:,i,:]), axis=1)
    pose_data = pose_data_bvh.reshape(pose_data_bvh.shape[0],-1,3)        
            
    for t in range(data.shape[0]):
        line = '%f %f %f ' % (data[t, 0], data[t, 1], data[t, 2])
        for d in range(pose_data.shape[1] - 1):
            line += '%f %f %f ' % (pose_data[t, d, 0], pose_data[t, d, 1], pose_data[t, d, 2])
        line += '%f %f %f\n' % (pose_data[t, -1, 0], pose_data[t, -1, 1], pose_data[t, -1, 2])
        fout.write(line)
    fout.close()


def plot_pose(pose, cur_frame, prefix, color='b'):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    
    parents = [-1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 12, 12, 12, 13, 14, 16, 17, 18, 19] # smpl
    # parents = [-1, 0, 1, 2, 3, 0, 5, 6, 7, 0, 9, 10, 11, 12, 11, 14, 15, 16, 11, 18, 19, 20]  # lafan
    # parents = [-1, 0, 1, 2, 3, 0, 5, 6, 7, 0, 9, 10, 11, 10, 13, 14, 15, 16, 16, 10, 19, 20, 21, 22, 22]  # h36m 25
    # parents = [-1, 0, 1, 2, 3, 4, 0, 6, 7, 8, 9, 0, 11, 12, 13, 14, 12, 16, 17, 18, 19, 20, 19, 22, 12, 24, 25, 26, 27, 28, 27, 30]  # h6m 32
    
    ax.cla()
    num_joint = pose.shape[0] // 3
    for i, p in enumerate(parents):
        # print(i, p)
        if i > 0:
            ax.plot([pose[i, 2], pose[p, 2]],
                    [pose[i, 0], pose[p, 0]],
                    [pose[i, 1], pose[p, 1]], c='r')

            ax.plot([pose[i + num_joint, 2], pose[p + num_joint, 2]],
                    [pose[i + num_joint, 0], pose[p + num_joint, 0]],  
                    [pose[i + num_joint, 1], pose[p + num_joint, 1]], c=color)

            ax.plot([pose[i + num_joint * 2, 2], pose[p + num_joint * 2, 2]],
                    [pose[i + num_joint * 2, 0], pose[p + num_joint * 2, 0]],
                    [pose[i + num_joint * 2, 1], pose[p + num_joint * 2, 1]], c='g')
    
    # Set the axis label 
    plt.xlabel("x")
    plt.ylabel("y")
    # plt.zlabel("z")
    
    plt.draw()
    plt.savefig(prefix + '_' + str(cur_frame) + '.png', dpi=200, bbox_inches='tight')
    plt.close('all')


def rearrange_joint_order(ori_data):
    ord_inv = [0, 9, 10, 11, 12, 13, 18, 19, 20, 21, 14, 15, 16, 17, 5, 6, 7, 8, 1, 2, 3, 4]
    arranged_data = torch.zeros(ori_data.shape).cuda()
    for i in range(ori_data.shape[2]):
        arranged_data[:, :, i, :] = ori_data[:, :, ord_inv[i], :]
    return arranged_data


def blend_quat(quat_list, iv_quat_list, root_list, iv_root_list, seq_len):
    quat = torch.cat(quat_list, 0)
    quat = quat.view(len(quat_list), -1, quat.shape[1], quat.shape[2]).permute(1, 0, 2, 3)
    quat = quat / torch.norm(quat, dim=-1, keepdim=True)

    # quat = rearrange_joint_order(quat)

    root_pos = torch.cat(root_list, 0)
    root_pos = root_pos.view(len(root_list), -1, root_pos.shape[1]).permute(1, 0, 2)

    iv_quat = torch.cat(iv_quat_list, 0)
    iv_quat = iv_quat.view(len(iv_quat_list), -1, iv_quat.shape[1], iv_quat.shape[2]).permute(1, 0, 2, 3)
    iv_quat = torch.flip(iv_quat, dims=[1])
    iv_quat = iv_quat / torch.norm(iv_quat, dim=-1, keepdim=True)

    # iv_quat = rearrange_joint_order(iv_quat)

    iv_root_pos = torch.cat(iv_root_list, 0)
    iv_root_pos = iv_root_pos.view(len(iv_root_list), -1, iv_root_pos.shape[1]).permute(1, 0, 2)
    iv_root_pos = torch.flip(iv_root_pos, dims=[1])

    bvh_list = []
    quat_list = []
    root_list = []

    over_len = (seq_len // 10) * 2
    half_len = seq_len // 2

    for i in range(0, seq_len):
        if i < (half_len - over_len):
            frame = torch.cat([root_pos[:, i], quat[:, i].view(quat.shape[0], -1)], -1)
            # frame = torch.cat([gt_root[:, i], quat[:, i].view(quat.shape[0], -1)], -1)
            quat_list.append(quat[:, i].view(quat.shape[0], -1))
            root_list.append(root_pos[:, i])
        elif i > (half_len + over_len - 1):
            index = i - half_len + over_len
            frame = torch.cat([iv_root_pos[:, index], iv_quat[:, index].view(iv_quat.shape[0], -1)], -1)
            # frame = torch.cat([gt_root[:, i], iv_quat[:, index].view(iv_quat.shape[0], -1)], -1)
            quat_list.append(iv_quat[:, index].view(iv_quat.shape[0], -1))
            root_list.append(iv_root_pos[:, index])
        else:
            ratio = float(i - half_len + over_len) / (over_len * 2 - 1)
            slerp_quat = quat_slerp(np.asarray(quat[:, i].cpu().detach().numpy()),
                                    np.asarray(iv_quat[:, i - half_len + over_len].cpu().detach().numpy()), ratio)
            slerp_quat = torch.from_numpy(slerp_quat)
            lerp_root = pos_lerp(root_pos[:, i], iv_root_pos[:, i - half_len + over_len], ratio)

            frame = torch.cat([lerp_root.cuda(), slerp_quat.cuda().view(slerp_quat.shape[0], -1)], -1)
            # frame = torch.cat([gt_root[:, i], slerp_quat.view(slerp_quat.shape[0], -1)], -1)
            quat_list.append(slerp_quat.cuda().view(slerp_quat.shape[0], -1))
            root_list.append(lerp_root)
        bvh_list.append(frame)

    pred_quat = torch.cat(quat_list, 0)
    pred_quat = pred_quat.view(len(quat_list), -1, 88).permute(1, 0, 2)
    pred_root = torch.cat(root_list, 0)
    pred_root = pred_root.view(len(root_list), -1, 3).permute(1, 0, 2)

    return bvh_list, pred_quat, pred_root


from scipy import signal
from scipy.spatial.transform import Rotation, Slerp

def l_slerp(forward_frame, backend_frame, root, backend_root, ratio):
    B, J, _ = forward_frame.shape
    blended_quat = torch.zeros((B, J, 4))
    for b in range(B):
        for j in range(J):
            forward = forward_frame[b, j]
            backend = backend_frame[b, j]

            keyframes = torch.cat([forward.unsqueeze(0), backend.unsqueeze(0)], 0)
            keyframes = keyframes.view(-1, 4)  # 2*4
            key_rots = Rotation.from_quat(keyframes.cpu())

            key_times = [0, 1]
            S = Slerp(key_times, key_rots)
            key_times.insert(1, ratio)
            interp_rots = S(key_times)

            blended_quat[b, j] = torch.from_numpy(interp_rots[1].as_quat())

    blended_root = root * (1 - ratio) + backend_root * ratio

    return blended_quat, blended_root

def original_blend_quat(quat_list, iv_quat_list, root_list, iv_root_list, gt_root, seq_len):
    quat = torch.cat(quat_list, 0)
    quat = quat.view(len(quat_list), -1, quat.shape[1], quat.shape[2]).permute(1, 0, 2, 3)

    quat = quat.view(quat.size(0), quat.size(1), -1).permute(0, 2, 1)  # B*88*T
    b, a = signal.butter(3, 0.2, 'lowpass')  # 配置滤波器 8 表示滤波器的阶数
    quat = signal.filtfilt(b, a, quat.cpu())
    quat = torch.from_numpy(np.ascontiguousarray(quat)).permute(0, 2, 1)
    quat = quat.view(quat.size(0), quat.size(1), -1, 4).cuda()
    quat = quat / torch.norm(quat, dim=-1, keepdim=True)

    root_pos = torch.cat(root_list, 0)
    root_pos = root_pos.view(len(root_list), -1, root_pos.shape[1]).permute(1, 0, 2)

    iv_quat = torch.cat(iv_quat_list, 0)
    iv_quat = iv_quat.view(len(iv_quat_list), -1, iv_quat.shape[1], iv_quat.shape[2]).permute(1, 0, 2, 3)
    iv_quat = torch.flip(iv_quat, dims=[1])

    iv_quat = iv_quat.view(iv_quat.size(0), iv_quat.size(1), -1).permute(0, 2, 1)  # B*88*T
    b, a = signal.butter(3, 0.2, 'lowpass')  # 配置滤波器 8 表示滤波器的阶数
    iv_quat = signal.filtfilt(b, a, iv_quat.cpu())
    iv_quat = torch.from_numpy(np.ascontiguousarray(iv_quat)).permute(0, 2, 1)
    iv_quat = iv_quat.view(iv_quat.size(0), iv_quat.size(1), -1, 4).cuda()
    iv_quat = iv_quat / torch.norm(iv_quat, dim=-1, keepdim=True)

    iv_root_pos = torch.cat(iv_root_list, 0)
    iv_root_pos = iv_root_pos.view(len(iv_root_list), -1, iv_root_pos.shape[1]).permute(1, 0, 2)
    iv_root_pos = torch.flip(iv_root_pos, dims=[1])

    bvh_list = []
    gq_list = []

    over_len = (seq_len // 10) * 2
    half_len = seq_len // 2

    for i in range(0, seq_len):
        if i < (half_len - over_len):
            frame = torch.cat([root_pos[:, i], quat[:, i].view(quat.shape[0], -1)], -1)
            # frame = torch.cat([gt_root[:, i], quat[:, i].view(quat.shape[0], -1)], -1)
            gq_list.append(quat[:, i].view(quat.shape[0], -1))
        elif i > (half_len + over_len - 1):
            index = i - half_len + over_len
            frame = torch.cat([iv_root_pos[:, index], iv_quat[:, index].view(iv_quat.shape[0], -1)], -1)
            # frame = torch.cat([gt_root[:, i], iv_quat[:, index].view(iv_quat.shape[0], -1)], -1)
            gq_list.append(iv_quat[:, index].view(iv_quat.shape[0], -1))
        else:
            ratio = float(i - half_len + over_len) / (over_len * 2 - 1)
            # slerp_quat = slerp(quat[:, i], iv_quat[:, i - half_len + over_len], ratio).cuda()
            slerp_quat, slerp_root = l_slerp(quat[:, i], iv_quat[:, i - half_len + over_len],
                                             root_pos[:, i], iv_root_pos[:, i - half_len + over_len],
                                             ratio)
            frame = torch.cat([slerp_root.cuda(), slerp_quat.cuda().view(slerp_quat.shape[0], -1)], -1)
            # frame = torch.cat([gt_root[:, i], slerp_quat.view(slerp_quat.shape[0], -1)], -1)
            gq_list.append(slerp_quat.cuda().view(slerp_quat.shape[0], -1))
        bvh_list.append(frame)

    pred_quat = torch.cat(gq_list, 0)
    pred_quat = pred_quat.view(len(gq_list), -1, 88).permute(1, 0, 2)

    return bvh_list, pred_quat